package com.tianji.aigc.memory;

import cn.hutool.core.collection.CollStreamUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.json.JSONUtil;
import com.tianji.common.utils.UserContext;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.util.List;

/**
 * @author Administrator
 * @version 1.0
 * @description: TODO
 * @date 2025/7/21 17:32
 */
public class RedisChatMemory implements ChatMemory {
    @Autowired
    private StringRedisTemplate redisTemplate;

    private String prefix = "CHAT:";

    public RedisChatMemory() {
    }

    public RedisChatMemory(String prefix,StringRedisTemplate redisTemplate) {
        this.prefix = prefix;
    }

    public RedisChatMemory(StringRedisTemplate redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    /**
     * 1、添加会话记忆
     * @param conversationId 会话ID
     * @param messages 会话消息
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        // 1、判断消息集合是否为空，为空，直接结束
        if(ObjectUtil.isEmpty(messages))return;

        // 2、获取key
        String key = this.getKey(conversationId);

        // 3、循环添加会话记忆
        redisTemplate.opsForList().rightPushAll(key, CollStreamUtil.toList(messages,message -> ChatMessage.toJson(message)));
    }

    /**
     * 2、获取会话记忆
     * @param conversationId 会话ID
     * @param lastN 获取的会话消息条数
     * @return
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        // 1、判断角标合理性
        if(lastN < 0){
            return List.of();
        }

        // 2、获取会话记忆
        String key = this.getKey(conversationId);
        List<String> chatMemoryJsonList = this.redisTemplate.opsForList().range(key, 0, lastN);
        if(ObjectUtil.isEmpty(chatMemoryJsonList)){
            return List.of();
        }

        // 3、转换为Message集合
        return CollStreamUtil.toList(chatMemoryJsonList,chatMemoryJson -> ChatMessage.toMessage(chatMemoryJson));
    }

    /**
     * 3、清除会话记忆
     * @param conversationId 会话ID
     */
    @Override
    public void clear(String conversationId) {
        String key = this.getKey(conversationId);
        this.redisTemplate.delete(key);
    }

    private String getKey(String conversationId) {
        return this.prefix + conversationId;
    }

    /**
     * 根据对话ID优化对话记录，删除最后的2条消息，因为这2条消息是从路由智能体存储的，请求由后续的智能体处理
     * 为了确保历史消息的完整性，所以需要将中间转发的消息清理掉
     *
     * @param conversationId 对话的唯一标识符
     */
    public void optimization(String conversationId) {
        this.redisTemplate.opsForList().rightPop(this.getKey(conversationId),2);
    }
}
